"""
FCN (Fully Convolutional Network) 模型实现
基于ResNet-50的FCN，使用跳跃连接
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet50, ResNet50_Weights


class FCN(nn.Module):
    """
    FCN模型（FCN-ResNet50）
    输入：(B, 3, H, W) - 需要3通道以使用ImageNet预训练
    输出：(B, out_channels, H, W)
    """

    def __init__(self, in_channels=3, out_channels=1, pretrained=True):
        """
        Args:
            in_channels: 输入通道数（必须是3）
            out_channels: 输出通道数
            pretrained: 是否使用预训练权重
        """
        super(FCN, self).__init__()

        if in_channels != 3:
            print(f"[WARNING] FCN 需要3通道输入以使用预训练权重，当前输入通道数: {in_channels}")

        # 使用ResNet-50作为backbone
        if pretrained:
            resnet = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
        else:
            resnet = resnet50(weights=None)

        # 编码器
        self.conv1 = resnet.conv1
        self.bn1 = resnet.bn1
        self.relu = resnet.relu
        self.maxpool = resnet.maxpool

        self.layer1 = resnet.layer1  # 256 channels, 1/4
        self.layer2 = resnet.layer2  # 512 channels, 1/8
        self.layer3 = resnet.layer3  # 1024 channels, 1/16
        self.layer4 = resnet.layer4  # 2048 channels, 1/32

        # FCN头部
        # 从layer4输出（2048通道）
        self.score_layer4 = nn.Conv2d(2048, out_channels, 1)

        # 从layer3输出（1024通道）
        self.score_layer3 = nn.Conv2d(1024, out_channels, 1)

        # 从layer2输出（512通道）
        self.score_layer2 = nn.Conv2d(512, out_channels, 1)

        # 上采样层
        self.upsample_2x = nn.ConvTranspose2d(out_channels, out_channels, 4, stride=2, padding=1)
        self.upsample_4x = nn.ConvTranspose2d(out_channels, out_channels, 4, stride=2, padding=1)
        self.upsample_8x = nn.ConvTranspose2d(out_channels, out_channels, 16, stride=8, padding=4)

    def forward(self, x):
        size = x.shape[2:]

        # 编码器
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)  # 1/4
        layer2_out = self.layer2(x)  # 1/8
        layer3_out = self.layer3(layer2_out)  # 1/16
        layer4_out = self.layer4(layer3_out)  # 1/32

        # FCN-8s：融合layer4, layer3, layer2
        # 从layer4开始
        score4 = self.score_layer4(layer4_out)
        score4_up = self.upsample_2x(score4)  # 上采样2倍到1/16

        # 加上layer3
        score3 = self.score_layer3(layer3_out)
        score3 = score3 + score4_up
        score3_up = self.upsample_4x(score3)  # 上采样2倍到1/8

        # 加上layer2
        score2 = self.score_layer2(layer2_out)
        score2 = score2 + score3_up

        # 最终上采样到原始尺寸
        out = self.upsample_8x(score2)

        # 裁剪到原始尺寸（如果有偏差）
        if out.shape[2:] != size:
            out = F.interpolate(out, size=size, mode='bilinear', align_corners=True)

        return out


if __name__ == '__main__':
    # 测试模型
    model = FCN(in_channels=3, out_channels=1, pretrained=False)
    x = torch.randn(2, 3, 256, 256)
    y = model(x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {y.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")
